import time

import mujoco
import mujoco.viewer
import numpy as np
from etils import epath

from mujoco_playground._src.locomotion.g1.base import get_assets

model = mujoco.MjModel.from_xml_string(epath.Path('/home/huang/mujoco_playground/mujoco_playground/_src/locomotion/g1/xmls/scene_mjx_feetonly_flat_terrain.xml').read_text(), assets=get_assets())
data  = mujoco.MjData(model)

traj = np.load('npz/g1/dance1_subject1.xpos.npz')        # 必须包含 qpos, qvel 两 key
qpos_traj = traj['qpos']          # shape (T, nq)
qvel_traj = traj['qvel']          # shape (T, nv)
geom_xpos_traj = traj['geom_xpos']          # shape (T, nv)
with mujoco.viewer.launch_passive(model, data) as viewer:
    for t in range(len(qpos_traj)):
        data.qpos = qpos_traj[t]
        data.qvel = qvel_traj[t]
        mujoco.mj_forward(model, data)   # 更新所有派生量
        # print(geom_xpos_traj[t][model.geom("left_foot").id][2])
        viewer.sync()
        # 控制回放速度，例如 1× 真实时间
        time.sleep(model.opt.timestep)